import os
import numpy as np
from sklearn.neighbors import KNeighborsClassifier
import matplotlib.pyplot as plt


# 处理数据
def build_data(dir_name):
    """
    构建数据
    :param dir_name: 对应文件夹的名称
    :return: 训练集或者测试集
    """
    # 遍历文件夹
    file_name_list = os.listdir(dir_name)

    data_arr = np.zeros((len(file_name_list), 1025))

    # 循环处理每一个文件
    for file_index, file_name in enumerate(file_name_list):
        # 加载文件里面的内容
        file_data = np.loadtxt(dir_name + "/" + file_name, dtype=np.str)
        # this is key
        single_arr = np.zeros((32, 32))
        for file_content_index, file_content in enumerate(file_data):
            # 转换为列表推导式
            # 处理成数值型的文件
            arr = [int(i) for i in file_content]
            # 方法二
            # 使用map 将file_content的每一个元素用int函数处理 然后转换为一个list
            # arr = list(map(int, file_content))
            # print(arr)
            single_arr[file_content_index, :] = arr

            # 展开
        data_arr[file_index, :1024] = single_arr.ravel()
        data_arr[file_index, 1024] = int(file_name.split('_')[0])
    # 返回完整的数据集
    return data_arr


def save_data(data, file_name):
    """
    保存数据
    :param data: 需要保存的数据集
    :param file_name: 保存的文件名
    :return: 文件路径
    """
    folder_name = './data/'
    if not os.path.exists(folder_name):
        os.makedirs(folder_name)
    np.save(folder_name + file_name, data)

    return folder_name + file_name


def load_data(training_path, test_path):
    """
    加载数据
    :param training_path: 训练集文件路径
    :param test_path: 测试集文件路径
    :return:
    """
    training = np.load(training_path)
    test = np.load(test_path)

    return training, test


def show_data(k_list, score_list, k_label, score_label, title):
    plt.figure()
    # 支持中文
    plt.rcParams['font.sans-serif'] = 'SimHei'
    # 支持负号
    plt.rcParams['axes.unicode_minus'] = False
    plt.plot(k_list, score_list)
    plt.title(title)
    plt.xlabel(k_label)
    plt.ylabel(score_label)
    for i, j in zip(k_list, score_list):
        plt.text(i, j, '%.2f' % j, horizontalalignment='center')

    plt.savefig('./' + title + '.png')
    plt.show()


def main():
    # 1.构建数据
    data_training = build_data('src/trainingDigits')
    data_test = build_data('src/testDigits')
    # 2.保存数据
    data_name_training = 'training.npy'
    data_name_test = 'test.npy'

    training_path = save_data(data_training, data_name_training)
    test_path = save_data(data_test, data_name_test)
    # 3.加载数据
    training, test = load_data(training_path, test_path)

    # 4.使用KNN算法
    k_list = [5, 6, 7, 8, 9, 10]
    score_list = []
    for k in k_list:
        # 实例化对象
        knn = KNeighborsClassifier(n_neighbors=k)
        # 训练数据
        knn.fit(training[:, :1024], training[:, 1024])
        # 预测数据

        # 计算准确率
        score = knn.score(test[:, :1024], test[:, 1024])
        print("当k = %d, 准确率为: %f" % (k, score))
        score_list.append(score)

    title = '不同k值时，准确率的变化趋势'
    k_label = 'k'
    score_label = 'score'
    show_data(k_list, score_list, k_label, score_label, title)


if __name__ == '__main__':
    main()
